package rbtree;

import avltree.AVLTree;

/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: a
 * Date: 2022-11-09
 * Time: 17:03
 */
public class RBTree {
    static class RBTreeNode{
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
        public int val;
        public COLOR color;

        public RBTreeNode(int val){
            this.val = val;
            this.color = COLOR.RED;
        }
    }

    public RBTreeNode root;

    public boolean insert(int val){
        RBTreeNode node = new RBTreeNode(val);
        if(root == null){
            root = node;
            root.color = COLOR.BLACK;
            return true;
        }
        //接下来的步骤类似于二叉搜索树的插入
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while(cur != null){
            if(cur.val < val){
                parent = cur;
                cur = cur.right;
            }else if(cur.val == val){
                return false;
            }else{
                parent = cur;
                cur = cur.left;
            }
        }
        if(parent.val < val){
            parent.right = node;
        }else{
            parent.left = node;
        }
        node.parent = parent;
        cur = node;
        while(parent != null && parent.color == COLOR.RED){
            RBTreeNode grandFather = parent.parent;
            if(parent == grandFather.left){
                RBTreeNode uncle = grandFather.right;
                if(uncle != null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    //继续向上修改
                    cur = grandFather;
                    parent = cur.parent;
                }else{
                    //uncle不存在 或者 uncle为黑
                    //情况三:
                    if(cur == parent.right){
                        rotateLeft(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }  //此时情况三变成情况二

                    //情况二:
                    rotateRight(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }else{
                //parent == grandFather.right
                RBTreeNode uncle = grandFather.left;
                if(uncle != null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    //继续向上修改
                    cur = grandFather;
                    parent = cur.parent;
                }else{
                    //uncle不存在 或者 uncle为黑
                    //情况三:
                    if(cur == parent.left){
                        rotateRight(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }  //此时情况三变成情况二

                    //情况二:
                    rotateLeft(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }
        }
        root.color = COLOR.BLACK;
        return true;
    }

    //左单旋
    private void rotateLeft(RBTreeNode parent){
        RBTreeNode subR = parent.right;
        RBTreeNode subRL = subR.left;
        RBTreeNode pParent = parent.parent;
        parent.right = subRL;
        //subRL存在的时候才执行
        if(subRL != null){
            subRL.parent = parent;
        }
        subR.left = parent;
        parent.parent = subR;

        if(parent == root){
            root = subR;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subR;
            }else{
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
    }

    //右单旋
    private void rotateRight(RBTreeNode parent){
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;
        RBTreeNode pParent = parent.parent;  //记录parent的父节点
        parent.left = subLR;
        //subLR存在的时候才执行
        if(subLR != null) {
            subLR.parent = parent;
        }
        subL.right = parent;
        parent.parent = subL;

        if(parent == root){
            subL = root;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subL;
            }else{
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
    }

    public boolean isRBTree(){
        if(root == null){
            return true;
        }
        if(root.color != COLOR.BLACK){
            System.out.println("根节点必须是黑色的");
            return false;
        }
        int blackNum = 0;
        RBTreeNode cur = root;
        while(cur != null){
            if(cur.color == COLOR.BLACK){
                blackNum++;
            }
            cur = cur.left;
        }
        return checkRedColor(root) && checkBlackNum(root, 0, blackNum);
    }

    private boolean checkRedColor(RBTreeNode root){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.RED){
            RBTreeNode parent = root.parent;
            if(parent.color == COLOR.RED){
                System.out.println("不能出现连续两个红色节点");
                return false;
            }
        }
        return checkRedColor(root.left) && checkRedColor(root.right);
    }

    private boolean checkBlackNum(RBTreeNode root, int pathBlackNum, int blackNum){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.BLACK){
            pathBlackNum++;
        }
        if(root.left == null && root.right == null){
            if(pathBlackNum != blackNum){
                System.out.println("路径上的黑色节点个数是不一样的");
                return false;
            }
        }
        return checkBlackNum(root.left, pathBlackNum, blackNum) && checkBlackNum(root.right, pathBlackNum, blackNum);
    }

    //中序遍历
    public void inorder(RBTreeNode root){
        if(root == null){
            return;
        }
        inorder(root.left);
        System.out.print(root.val + " ");
        inorder(root.right);
    }
}
